import torch
import torchvision.transforms as transforms
from model import ContactAngle
from config import *
from PIL import Image
import os
import datetime
from train import args, test_data_loader, Loss

# 使用CPU或者GPU进行预测
if torch.cuda.is_available():
    device = torch.device('cuda')
else:
    device = torch.device('cpu')

# 图像预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])

start = datetime.datetime.now()
# 模型加载
model = ContactAngle(model_name=args.mn, out_puts=1020)
model.load_state_dict(torch.load(weight_root + r'\model-%s.pth' % args.mn))
model = model.to(device)
end = datetime.datetime.now()
print('预测模型加载时间:%.2f' % (end - start).seconds)

# 在测试集上计算模型预测接触角的平均误差
def mean_error_value():
    model.eval()
    with torch.no_grad():
        test_loss = 0
        for step, (img_a, txt_a, label_a) in enumerate(test_data_loader):
            img_b = img_a.to(device)
            txt_b = txt_a.to(device)
            label = label_a.to(device)
            outputs = model(img_b, txt_b)
            loss = Loss(outputs, label)
            test_loss += loss.item()
        test_loss /= len(test_data_loader)
        print('The Testing Loss is:%.4f' % test_loss)


# 通过用户输入参数进行接触角预测
num = 1
while True:
    print('\n第%d次预测' % num)
    print('----------------------------------------\n')
    txt = []

    out1 = input('请输入打标次数(0-200):')
    if out1.isdigit() and 0 <= int(out1) <= 200:
        x1 = float(out1)
    else:
        print('请输入一个0~200的整数！！！')
        continue

    out2 = input('请输入激光功率(0-100%):')
    if out2.isdigit() and 0 <= int(out2) <= 100:
        x2 = float(out2)
    else:
        print('请输入一个0~100的整数！！！')
        continue

    out3 = input('请输入打标速度(0-600):')
    if out3.isdigit() and 0 <= int(out3) <= 600:
        x3 = float(out3)
    else:
        print('请输入一个0~600的整数！！！')
        continue

    out4 = input('请输入沟槽间距(0-0.25):')
    if out4.count('.') == 1 and not out4.startswith('.') and not out4.endswith('.') and 0 <= float(out4) <= 0.25:
        x4 = float(out4)
    else:
        print('请输入一个范围在0~0.25的小数！！！')
        continue

    root = input('请输入图像路径：')
    if not os.path.exists(root):
        print('请输入正确图像路径!!!')
        print('参考路径如下：‘E:/image.jpg')
        continue
    else:
        img = Image.open(root)
        img = transform(img).unsqueeze(0)
    txt = torch.Tensor([x1 / 200, x2 / 100, x3 / 600, x4 / 0.25]).unsqueeze(0)

    model.eval()
    with torch.no_grad():
        img = img.to(device)
        txt = txt.to(device)
        predict = model(img, txt)
        print('预测的接触角大小为:%.4f' % (predict.item() * 180))
        num += 1
        quit_a = input('若要退出程序，请按q!:')
        print('\n----------------------------------------\n')
        if quit_a == 'q':
            break
        else:
            continue
